'''
breif:
    retrieval demo
    This demo is based on the GLM-4-flash model
    The demo provied user upload file and make their own knowledge base which can be used to model.
'''
from zhipuai import ZhipuAI
import json
content  =  open('API.json','r',encoding='utf-8').read()
api = json.loads(content)['api']
client = ZhipuAI(api_key=api)

class AI():
    def __init__(self):
        ...
    def upload(self,file_path,knowledge_id,sentence_size=202):
        resp = client.knowledge.document.create(
            file=open(file_path, "rb"),
            purpose="retrieval",
            knowledge_id=knowledge_id,
            sentence_size=sentence_size,
            custom_separator=["\n"]
        )
        return resp
    def creat_knowledge(self,name,description="Null"):
        result = client.knowledge.create(
            embedding_id=3,
            name=name,
            description=description
        )
        return result.id
    def chat(self,text):
        response = client.chat.completions.create(
            model="glm-4-flash",  # 填写需要调用的模型名称
            messages=[
                {"role": "user", "content": text},
            ],
            tools=[
                    {
                        "type": "retrieval",
                        "retrieval": {
                            "knowledge_id": "1873608325132926976",
                            "prompt_template": "从文档\n\"\"\"\n{{knowledge}}\n\"\"\"\n中找问题\n\"\"\"\n{{question}}\n\"\"\"\n的答案，找到答案就仅使用文档语句回答问题，找不到答案就用自身知识回答并且告诉用户该信息不是来自文档。\n不要复述问题，直接开始回答。"
                        }
                    }
                    ],
            stream=True,
        )
        for trunk in response:
            content = content = trunk.choices[0].delta.content
            print(content, end="")